################################################################################
#
# SoftwareScanner
#
# Create table for scan hashes
# +----------+-----------+----------+----------+------------+------------+
# | ID       | SHA1      | MD5      | Count    | N          | Q          |
# +----------+-----------+----------+----------+------------+------------+
#
################################################################################

import pyflag.DB as DB
import pyflag.FlagFramework as FlagFramework
import pyflag.pyflaglog as pyflaglog
import pyflag.conf
from SoftwareUtil import getProductCount
from pyflag.Scanner import *
from hashlib import md5, sha1


config = pyflag.conf.ConfObject()
config.add_option("SOFT_SCANNER_HASHES", default = "SoftScannerHashes")



class SoftwareScannerTables(FlagFramework.EventHandler):
    def create(self, dbh, case):
        dbh.execute("""CREATE TABLE IF NOT EXISTS %s (
            `ID` int(32) NOT NULL auto_increment,
            `SHA1` char(40) NOT NULL,
            `MD5` char(32) NOT NULL,
            `Count` int(16) NOT NULL default '1',
            `N` int(16) NOT NULL default '0',
            PRIMARY KEY (`ID`),
            KEY sha1_index (`SHA1`))""", config.SOFT_SCANNER_HASHES)



class HashScanFactory(GenScanFactory):
    """Hash Scanner [Counting SHA-1 & MD5 for files]"""
    default = True
    #depends = ['TypeScan']
    group = 'GeneralForensics'

    def __init__(self, fsfd):
        GenScanFactory.__init__(self, fsfd)


    class Scan(ScanIfType):
        def __init__(self, inode, ddfs, outer, factories = None, fd = None):
            ScanIfType.__init__(self, inode, ddfs, outer, factories, fd = fd)
	    path, inode, inode_id = ddfs.lookup(inode = fd.inode)
	    pyflaglog.log(pyflaglog.INFO, "Hash scanner: scanning file %s" % path)
	    ## Get MD5 & SHA-1 object
            self.md5Obj = md5()
	    self.sha1Obj = sha1()


        def boring(self, metadata, data = None):
	    return False


        def process(self, data, metadata = None):
	    ScanIfType.process(self, data, metadata)
	    self.md5Obj.update(data)
    	    self.sha1Obj.update(data)


        def finish(self):
	    #ScanIfType.finish(self)
            md5 = self.md5Obj.hexdigest().upper()
            sha1 = self.sha1Obj.hexdigest().upper()
            pyflaglog.log(pyflaglog.INFO, '[SHA-1: %s] [MD5: %s]' % (sha1, md5))

            try:
                dbh = DB.DBO(self.case)
                dbh.execute("""SELECT ID, Count FROM %s WHERE SHA1 = '%s' AND MD5 = '%s'""",
                    (config.SOFT_SCANNER_HASHES, sha1, md5))
                row = dbh.fetch()

                if row == None:
                    n = getProductCount(sha1, md5)
                    dbh.insert(config.SOFT_SCANNER_HASHES,
                        SHA1 = sha1,
                        MD5 = md5,
                        Count = 1,
                        N = n,
                        )
                else:
                    id = row['ID']
                    count = int(row['Count']); count += 1;
                    dbh.update(config.SOFT_SCANNER_HASHES,
                        "ID = %s" % id,
                        Count = count,
                        )
            except DB.DBError, e:
                pyflaglog.log(pyflaglog.ERROR, str(e))
            finally:
                pass

